#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#include "rsa.h"

int mp[256];

void init() {
	srand(time(0));
	memset(mp, -1, sizeof mp);
	for (char i = '0'; i <= '9'; ++i) {
		mp[(int)i] = (int)(i - '0');
	}
	for (char i = 'a'; i <= 'f'; ++i) {
		mp[(int)i] = 10 + (int)(i - 'a');
	}
	for (char i = 'A'; i <= 'F'; ++i) {
		mp[(int)i] = 10 + (int)(i - 'A');
	}
}

int read_big_integer(big_integer *x, FILE *fp) {
	static char buffer[K_BYTES * 2];
	char ch;
	int res;
	while (1) {
		res = fscanf(fp, "%c", &ch);
		if (res != 1) {
			return 1;
		}
		if (mp[(int)ch] != -1) {
			break;
		}
	}
	buffer[0] = ch;
	int idx = 1;
	while (1) {
		res = fscanf(fp, "%c", &ch);
		if (res != 1 || mp[(int)ch] == -1) {
			break;
		}
		buffer[idx++] = ch;
		if (idx > sizeof(buffer) / sizeof(buffer[0])) {
			return 1;
		}
	}

	clear_big_integer(x);
	uint64_t num;
	for (int i = 0; i < NUM_DIGITS; ++i) {
		for (int j = 0; j < 8; ++j) {
			num = (uint64_t)mp[(int)buffer[--idx]];
			x->digits[i] |= (num << (j * 4));
			if (idx == 0) {
				return 0;
			}
		}
	}
	return 0;
}

int load_keys(big_integer *n, big_integer *e, big_integer *d, FILE *fp) {
	int res;
	if ((res = read_big_integer(n, fp)) != 0) {
		return res;
	}
	if ((res = read_big_integer(e, fp)) != 0) {
		return res;
	}
	if ((res = read_big_integer(d, fp)) != 0) {
		return res;
	}
	return 0;
}

int save_keys(const big_integer *n, const big_integer *e, const big_integer *d, FILE *fp) {
	fprint_big_integer(n, fp);
	fprint_big_integer(e, fp);
	fprint_big_integer(d, fp);
	return 0;
}

int main() {
	init();

	/* 获得公私钥 */
	big_integer n, e, d;
	FILE *fp = fopen("keys.txt", "r");
	if (fp) {
		int res = load_keys(&n, &e, &d, fp);
		if (res != 0) {
			return res;
		}
	}
	else {
		generate_keys(&n, &e, &d);
		fp = fopen("keys.txt", "w");
		if (!fp) {
			return 2;
		}
		int res = save_keys(&n, &e, &d, fp);
		if (res != 0) {
			return res;
		}
	}
	fclose(fp);

	/* 简单测试加密、解密 */
	octet_string encrypted;
	char hello[] = "Hello. My name is Zhang3 and I study in SYSU.";
	printf("Message to encrypt: %s\n", hello);
	printf("Length: %d\n", (int)(sizeof(hello) / sizeof(hello[0])));
	
	printf("\n");
	if (rsaes_pkcs1_v15_encrypt(&encrypted, &n, &e, (uint8_t*)hello, sizeof(hello) / sizeof(hello[0])) != 0) {
		printf("Fail encrypting\n");
		return 1;
	}
	printf("Encrypted octet string in hex: ");
	print_octet_string(&encrypted);

	uint8_t *decrypted;
	size_t len;
	printf("\n");

	if (rsaes_pkcs1_v15_decrypt(&decrypted, &len, &n, &d, &encrypted)) {
		printf("Fail decrypting\n");
		return 1;
	}
	printf("Decrypted Message: %s\n", (char*)decrypted);
	printf("Length: %d\n", (int)len);

	free(decrypted);
	return 0;
}
